#pragma once
#include <vector>
#include "common\define.h"
#include "NvInfer.h"
#include "opencv2/opencv.hpp"
#include "cuda_img_resize.cuh"
#include "../txr_algo_dlm_seg/txr_algo_dlm_seg_define.h"

using namespace txr_algo_dlm_seg;
namespace yolo_v6
{
	struct DetectRes {
		int classes;
		float x;
		float y;
		float w;
		float h;
		float x1;
		float x2;
		float y1;
		float y2;
		float prob;
		float cof[32];
	};

	class YoloV6
	{
	public:
		YoloV6();
		~YoloV6();
		bool Init(int gpu_id, const char* dats_path);
		void Detect(st_detect_unit* p_unit, int num);
	protected:
	private:
		void Reset();
		bool LoadEngine(std::vector<char> v_engine_data);
		std::vector<std::vector<DetectRes>> PostProcess(st_img_rgb* p_imgs, int num, float* output, const int& outSize);
		void NmsDetect(std::vector<DetectRes>& detections, float nms_threshold);
		float IOUCalculate(const DetectRes& det_a, const DetectRes& det_b);

		std::vector<float> prepareImage(std::vector<cv::Mat> vec_Mat);
		std::vector<std::vector<DetectRes>> postProcess(
			const std::vector<cv::Mat>& vec_Mat, st_detect_unit* p_imgs, float* output, const int& outSize);
		std::vector<cv::Mat> postProcessMask(
			const std::vector<cv::Mat>& vec_Mat, st_detect_unit* p_imgs, float* output, const int& outSize, std::vector<std::vector<DetectRes>> vv_boxes);

	private:
		st_encrypt_info m_info;
		nvinfer1::ICudaEngine* m_engine = nullptr;
		nvinfer1::IExecutionContext* m_context = nullptr;
		std::vector<void*> m_v_dev_buffer;
		std::vector<int64_t> m_v_dev_buf_size;

		std::vector<std::vector<int>> m_grids;

		int m_out_size[2] = { 0 };
		cudaStream_t m_stream;
		std::vector<float> m_v_out_buf[2];

		std::vector<st_cuda_resize_dev_space> _v_resize_space;

		float m_img_mean[3] = { 0,0,0 };
		float m_img_std[3] = { 1,1,1 };
		float m_num_anchors[3] = { 3,3,3 };
		int m_num_rows = 0;

	};

}//namespace yolo_v6